from flask import render_template, request, session, redirect, url_for, Blueprint, jsonify, make_response, g
from flask_socketio import emit

from decorator import login_required
from exts import socketio
from form import LoginForm, RegisterForm, ProfileForm
from data_process_answer.answer import generate_response
from model import Usermodel, db, Conversation
from data_process_answer.vector_store import perform_query
from data_process_answer.answer import prompt_augmented

bp_main = Blueprint('main', __name__, url_prefix='/')

@bp_main.route('/register', methods=['GET', 'POST'])
def register():
    form = RegisterForm(request.form)
    if request.method == 'POST' and form.validate():
        username = form.username.data
        email = form.email.data
        password = form.password.data
        user = Usermodel(username=username, email=email, password=password)
        db.session.add(user)
        db.session.commit()
        # 注册成功后，重定向到登录页面，并携带注册成功的标志
        return redirect(url_for("main.login", registered=True))
    else:
        # 如果表单验证失败，打印错误信息但不重定向
        return render_template("register.html", form=form, errors=form.errors)

@bp_main.route('/check-email', methods=['POST'])
def check_email():
    email = request.form.get('email')
    user = Usermodel.query.filter_by(email=email).first()
    if user:
        return jsonify({'error': '该邮箱已经被注册！'})
    return jsonify({'success': True})

@bp_main.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'GET':
        return render_template("login.html")
    else:
        form = LoginForm(request.form)
        if form.validate():
            email = form.email.data
            password = form.password.data
            user = Usermodel.query.filter_by(email=email).first()
            if user and user.password == password:
                session['user_id'] = user.id

                # 设置 cookie
                response = make_response(redirect(url_for("main.chat")))
                response.set_cookie('user_id', str(user.id), max_age=30 * 24 * 60 * 60)  # 保存30天
                return response
            else:
                return redirect(url_for("main.login"))
        else:
            print(form.errors)
            return redirect(url_for("main.register"))


@bp_main.route('/get_chat_history/<int:user_id>', methods=['GET'])
@login_required
def get_chat_history(user_id):
    # 获取指定用户的聊天记录
    conversations = Conversation.query.filter_by(user_id=user_id).all()
    chat_history = [
        {'user_input': conv.user_input, 'model_response': conv.model_response, 'timestamp': conv.create_time}
        for conv in conversations
    ]
    return jsonify(chat_history)

@bp_main.route('/logout')
# @login_required
def logout():
    session.pop('user_id', None)
    return redirect(url_for("main.login"))


@bp_main.route('/chat', methods=["GET", 'POST'])
# @login_required
def chat():
    if request.method == 'GET' and g.user is not None:
        return render_template("chat.html")
    elif g.user is None:
        return redirect(url_for("main.login"))
    else:
        user_input = request.json.get('message')
        context = perform_query(user_input, 3)
        prompt_aug = prompt_augmented(user_input, context)
        response = generate_response(prompt_aug)  # 确保这个函数存在并返回正确的响应
        conversation = Conversation(user_input=user_input, model_response=response, user_id=session.get('user_id'))
        db.session.add(conversation)
        db.session.commit()
        return jsonify({'response': response})  # 返回JSON格式的响应


@socketio.on('message')
def handle_message(message):
    response = generate_response(message)
    emit('response', {'data': response})


@bp_main.route('/profile', methods=['GET', 'POST'])
@login_required
def profile():
    form = ProfileForm(request.form)
    current_user = Usermodel.query.filter_by(id=session.get('user_id')).first()
    if current_user is None:
        print('用户信息未找到，请先注册或联系管理员。')
        return redirect(url_for('main.chat'))  # 重定向到首页或其他页面
    if request.method == 'POST' and form.validate():
        # 更新用户信息
        print(current_user)
        print(form.bio.data)
        current_user.bio = form.bio.data
        try:
            db.session.commit()
            print('个人信息已更新！')
        except Exception as e:
            db.session.rollback()
            print(f'更新失败：{str(e)}')
        return redirect(url_for('main.profile'))
    else:
        print(form.errors)
    return render_template('profile.html', form=form, user=current_user)


